##########################################################################################

library(data.table)
library(optparse)
library(ArchR)
library(Seurat)
library(ggsci)
library(dplyr)
library(GenomicFeatures)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--rna_data_file"), type = "character"),
    #make_option(c("--nclust"), type = "character"),
    make_option(c("--gtf_file"), type = "character"),
    make_option(c("--gff3_file"), type = "character"),
    make_option(c("--gene_type"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined_peak.combineRNA.qc.Rdata"

    ## 表达文件
    rna_data_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined.annotationCellType.qc.Rdata"

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    ## gtf_file 
    gtf_file <- "~/ref/GTF/20240317_gencode_v32_gene.bed"
    gff3_file <- "~/ref/GTF/gencode.v32.annotation.gff3"

    ## 
    nclust <- 11

    gene_type <- "protein_coding"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/peak2gene/germ/protein_coding"

}


###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
rna_data_file <- opt$rna_data_file
scriptPath <- opt$scriptPath
#nclust <- as.numeric(opt$nclust)
gene_type <- opt$gene_type
out_path <- opt$out_path
gtf_file <- opt$gtf_file
gff3_file <- opt$gff3_file

dir.create( out_path , recursive = T)

##########################################################################################
## 导入数据
a <- load(comine_data_file)
# testis_combined_peak_combineRNA
b <- load(rna_data_file)

## 细胞顺序
cell_order <- c("SSC", "Differenting&Differented SPG", "Leptotene",
    "Zygotene", "Patchytene", "Diplotene",
    "Early stage of spermatids", "Round&ElongateS.tids", "Sperm",
    "Leydig cells", "Myoid cells", "Pericytes",
    "Sertoli cells", "Endothelial cells", "NKT cells", "Macrophages"
    )

## 细胞颜色
use_colors <- c(pal_npg("nrc")(10) , pal_jco("default")(6))
names(use_colors) <- c("Myoid cells" , "Leydig cells" , "Endothelial cells" , "Zygotene" , "Round&ElongateS.tids" , 
"Patchytene" , "SSC" , "Sperm" , "Diplotene" , "Early stage of spermatids" , "Leptotene" , 
"Sertoli cells" , "Macrophages" , "Differenting&Differented SPG" , "Pericytes" , "NKT cells" )

##########################################################################################
## 参考的gtf文件
dat_gtf <- fread(gtf_file)
txdb <- makeTxDbFromGFF(gff3_file)

## 提取关系的基因集合（编码基因或lcnRNA）
if( gene_type == "protein_coding2lncRNA" ){
  dat_gtf <- subset( dat_gtf , V5 %in% c("protein_coding" , "lncRNA") )
}else{
  dat_gtf <- subset( dat_gtf , V5 == gene_type )
}

################################################################################
## 构建基因的位置信息
gene_ensg <- dat_gtf[,c("V6","V4")]
colnames(gene_ensg) <- c("GENEID" , "name")

gene_gr <- genes(txdb, columns = c("GENEID"))

gr_df <- data.frame(
  seqnames = seqnames(gene_gr),
  start = start(gene_gr),
  end = end(gene_gr),
  strand = strand(gene_gr),
  GENEID = names(mcols(gene_gr)$GENEID)
)

## gene name对应多个ensg的，选第一个
gr_df <- merge( gr_df , gene_ensg , by = "GENEID" )
gr_df <- gr_df %>% 
group_by( name ) %>%
summarize( seqnames = seqnames[1] , start = start[1] , end = end[1] , strand = "*"  )

gr <- GRanges(
  seqnames = Rle(gr_df$seqnames),
  ranges = IRanges(start = gr_df$start, end = gr_df$end),
  strand = gr_df$strand,
  name = gr_df$name
)

names(gr) <- gr_df$name

## 保留感兴趣的基因
use_gene <- rownames(scrnat@assays$RNA$counts)[rownames(scrnat@assays$RNA$counts) %in% names(gr)]
gr <- gr[use_gene]

## 提取RNA的counts矩阵
seruat_cds_all <- SummarizedExperiment(scrnat@assays$RNA$counts[use_gene,] , rowRanges = gr)
colnames(seruat_cds_all) <- gsub( "_" , "#" ,  colnames(seruat_cds_all) )

## 替换表达矩阵
projHeme5 <- testis_combined_peak_combineRNA
projHeme5 <- addGeneExpressionMatrix(input = projHeme5, seRNA = seruat_cds_all, force = TRUE)
projHeme5 <- addImputeWeights(projHeme5)

#motifMatrix <- getMatrixFromProject(projHeme5, useMatrix="GeneExpressionMatrix")

##########################################################################################
## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

##########################################################################################
## To identify peak-to-gene links in ArchR
projHeme5 <- addPeak2GeneLinks(ArchRProj = projHeme5 , useMatrix = "GeneExpressionMatrix")

## 参数参考
## https://github.com/GreenleafLab/scScalpChromatin/blob/9e333bd3194ed6548cfad35c4b9ba678b0cdde31/Figure_2_Linked_Peaks.R#L71
# P2G definition cutoffs
corrCutoff <- 0.5       # Default in plotPeak2GeneHeatmap is 0.45
varCutoffATAC <- 0.25   # Default in plotPeak2GeneHeatmap is 0.25
varCutoffRNA <- 0.25    # Default in plotPeak2GeneHeatmap is 0.25


for( nclust in 5:20 ){
  print(nclust)

  ########################################
  ## Plotting a heatmap of peak-to-gene links
  #projHeme5@cellColData$cell_type <- factor( projHeme5@cellColData$cell_type , levels = cell_order , order = T )
  #nclust <- 25 
  p <- plotPeak2GeneHeatmap(
    projHeme5, 
    corCutOff = corrCutoff, 
    groupBy="cell_type", 
    nPlot = 1000000, returnMatrices=FALSE, 
    k=nclust, seed=1, palGroup=use_colors
    )

  out_file <- paste0(out_path, sprintf("/peakToGeneHeatmap_LabelClust_k%s.pdf", nclust))
  pdf(out_file , width=16, height=15)
  draw(p)
  dev.off()

  ########################################
  ## 计算每个cluster里面peak富集在哪些motif
  atac_proj <- projHeme5

  # Get all peaks
  allPeaksGR <- getPeakSet(atac_proj)
  allPeaksGR$peakName <- (allPeaksGR %>% {paste0(seqnames(.), "_", start(.), "_", end(.))})
  names(allPeaksGR) <- allPeaksGR$peakName

  # Need to force it to plot all peaks if you want to match the labeling when you 'returnMatrices'.
  p2gMat <- plotPeak2GeneHeatmap(
    atac_proj, 
    corCutOff = corrCutoff, 
    groupBy="cell_type",
    nPlot = 1000000, returnMatrices=TRUE, 
    k=nclust, seed=1)

  ## 输出peak-gene相关参数
  kclust_df_out <- cbind( kclust=p2gMat$ATAC$kmeansId , p2gMat$Peak2GeneLinks )
  tmp_peak_dat <- data.frame(atac_proj@peakSet)
  tmp_peak_dat$peak <- paste0( tmp_peak_dat$seqnames , ":" , tmp_peak_dat$start , "-" , tmp_peak_dat$end )
  kclust_df_out <- merge( kclust_df_out , tmp_peak_dat , by = "peak" )

  out_file <- paste0(out_path, sprintf("/peakToGeneHeatmap_LabelClust_k%s.tsv", nclust))
  write.table(kclust_df_out , out_file , row.names = F , sep = "\t")

  ## 记录每个基因对应多少peak在各自的clust里面
  peak_num <- data.frame(kclust_df_out) %>% 
  group_by( gene , kclust ) %>%
  summarize( peak_num = length(peak) )
  out_file <- paste0(out_path, sprintf("/peakToGeneHeatmap_LabelClust_k%s.peakNum.tsv", nclust))
  write.table(peak_num , out_file , row.names = F , sep = "\t")

  # Get association of peaks to clusters
  kclust_df <- data.frame(
    kclust=p2gMat$ATAC$kmeansId,
    peakName=p2gMat$Peak2GeneLinks$peak,
    gene=p2gMat$Peak2GeneLinks$gene
    )

  # Fix peakname
  kclust_df$peakName <- sapply(kclust_df$peakName, function(x) strsplit(x, ":|-")[[1]] %>% paste(.,collapse="_"))

  # Get motif matches
  matches <- getMatches(atac_proj, "Motif")
  r1 <- SummarizedExperiment::rowRanges(matches)
  rownames(matches) <- paste(seqnames(r1),start(r1),end(r1),sep="_")
  matches <- matches[names(allPeaksGR)]

  clusters <- unique(kclust_df$kclust) %>% sort()

  enrichList <- lapply(clusters, function(x){
    cPeaks <- kclust_df[kclust_df$kclust == x,]$peakName %>% unique()
    ArchR:::.computeEnrichment(matches, which(names(allPeaksGR) %in% cPeaks), seq_len(nrow(matches)))
    }) %>% SimpleList
  names(enrichList) <- clusters

  # Format output to match ArchR's enrichment output
  assays <- lapply(seq_len(ncol(enrichList[[1]])), function(x){
      d <- lapply(seq_along(enrichList), function(y){
          enrichList[[y]][colnames(matches),x,drop=FALSE]
        }) %>% Reduce("cbind",.)
      colnames(d) <- names(enrichList)
      d
    }) %>% SimpleList
  names(assays) <- colnames(enrichList[[1]])
  assays <- rev(assays)
  res <- SummarizedExperiment::SummarizedExperiment(assays=assays)

  formatEnrichMat <- function(mat, topN, minSig, clustCols=TRUE){
    plotFactors <- lapply(colnames(mat), function(x){
      ord <- mat[order(mat[,x], decreasing=TRUE),]
      ord <- ord[ord[,x]>minSig,]
      rownames(head(ord, n=topN))
    }) %>% do.call(c,.) %>% unique()
    pMat <- mat[plotFactors,]
    prettyOrderMat(pMat, clusterCols=clustCols)$mat
  }

  pMat <- formatEnrichMat(assays(res)$mlog10Padj, 5, 10, clustCols=FALSE)
  # Save maximum enrichment
  tfs <- strsplit(rownames(pMat), "_") %>% sapply(., `[`, 1)
  rownames(pMat) <- paste0(tfs, " (", apply(pMat, 1, function(x) floor(max(x))), ")")

  pMat <- apply(pMat, 1, function(x) x/max(x)) %>% t()
  rownames(pMat) <- sapply( strsplit(rownames(pMat) , " ") , "[" , 1)

  out_file <- paste0(out_path, sprintf("/enrichedMotifs_kclust_p2gHM_k%s.pdf", nclust))
  pdf(out_file ,width=12, height=15)
  ht_opt$simple_anno_size <- unit(0.25, "cm")
  hm <- BORHeatmap(
    pMat, 
    limits=c(0,1), 
    clusterCols=FALSE, clusterRows=FALSE,
    labelCols=TRUE, labelRows=TRUE,
    dataColors = cmaps_BOR$comet,
    #top_annotation = ta,
    row_names_side = "left",
    width = ncol(pMat)*unit(0.5, "cm"),
    height = nrow(pMat)*unit(0.4, "cm"),
    border_gp=gpar(col="black"), # Add a black border to entire heatmap
    legendTitle="Norm.Enrichment -log10(P-adj)[0-Max]"
    )
  draw(hm)
  dev.off()

  out_file <- paste0(out_path, sprintf("/enrichedMotifs_kclust_p2gHM_k%s.tsv", nclust))
  write.table(pMat , out_file , row.names = T , sep = "\t")


  if(gene_type!='lncRNA'){
    ########################################
    # GO enrichments of top N genes per cluster 
    # ("Top" genes are defined as having the most peak-to-gene links)
    kclust <- unique(kclust_df$kclust) %>% sort()
    all_genes <- kclust_df$gene %>% unique() %>% sort()

    # Save table of top linked genes per kclust
    nGOgenes <- 200
    topKclustGenes <- lapply(kclust, function(k){
      kclust_df[kclust_df$kclust == k,]$gene %>% getFreqs() %>% head(nGOgenes) %>% names()
      }) %>% do.call(cbind,.)
    outfile <- paste0(out_path, sprintf("/topN_genes_kclust_k%s.tsv", nclust))
    write.table(topKclustGenes, file=outfile, quote=FALSE, sep='\t', row.names = FALSE, col.names=TRUE)

    # 通路富集
    GOresults <- lapply(kclust, function(k){
      message(sprintf("Running GO enrichments on k cluster %s...", k))
      clust_genes <- topKclustGenes[,k]
      upGO <- rbind(
        calcTopGo(all_genes, interestingGenes=clust_genes, nodeSize=5, ontology="BP") 
        )

      upGO <- upGO[order(as.numeric(upGO$pvalue), decreasing=FALSE),]

      ## 构造GO对象
      geneList <- factor(as.integer(all_genes %in% clust_genes))
      names(geneList) <- all_genes
      GOdata <- suppressMessages(new(
        "topGOdata",
        ontology = "BP",
        allGenes = geneList,
        annot = annFUN.org, mapping = "org.Hs.eg.db", ID = "symbol",
        nodeSize = 5
      ))
      ## 提取每个通路中感兴趣的基因
      gene_in_pathway <- sapply(upGO$GO.ID, function(x)
        {
          genes <- genesInTerm(GOdata, x)
          # myGenes is the queried gene list
          paste0(genes[[1]][genes[[1]] %in% clust_genes] , collapse = "," )
        })
      
      upGO$gene_in_pathway <- gene_in_pathway
      upGO$cluster <- paste0("cluster" , k)

      return(upGO)
    })

    names(GOresults) <- paste0("cluster_", kclust)

    ## 输出为表格
    pathway_res <- c()
    for( clus in names(GOresults) ){

      tmp <- GOresults[[clus]]
      tmp$cluster <- clus
      pathway_res <- rbind( pathway_res , tmp )
    }

    out_file <- paste0(out_path, sprintf("/kclust_GO_3termsBPonlyBarLim_k%s.tsv", nclust))
    write.table( pathway_res , out_file , row.names = F , quote = F , sep ="\t" )


    # Plots of GO term enrichments:
    pdf(paste0(out_path, sprintf("/kclust_GO_3termsBPonlyBarLim_k%s.pdf", nclust)), width=10, height=6)
    for(name in names(GOresults)){
        goRes <- GOresults[[name]]
        if(nrow(goRes)>1){
          print(topGObarPlot(goRes, cmap = cmaps_BOR$comet, 
            nterms=10, border_color="black", 
            barwidth=0.85, title=name, barLimits=c(0, 15)))
        }
    }
    dev.off()
  }

}